// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.


#include "bfcallocator/allocatormgr.h"
#include "bfcallocator/bfc_arena.h"
#include <mutex>
#include <sstream>
#include <unordered_map>
#include <limits>
#include "bfcallocator/cuda_allocator.h"
#include "bfcallocator/cuda_call.h"

namespace bfcallocator {

namespace {
//It assumes max(OrtMemType) <= 1, min(OrtMemType) = -2
inline int MakeKey(int id, OrtMemType mem_type) {
  return id << 2 | (mem_type + 2);
}
}  // namespace

AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) {
  auto device_allocator = std::unique_ptr<IAllocator>(info.device_alloc_factory(info.device_id));

  if (info.use_arena) {
    size_t max_mem = info.arena_cfg.max_mem == 0 ? BFCArena::DEFAULT_MAX_MEM : info.arena_cfg.max_mem;
    int initial_chunk_size_bytes = info.arena_cfg.initial_chunk_size_bytes == -1
                                       ? BFCArena::DEFAULT_INITIAL_CHUNK_SIZE_BYTES
                                       : info.arena_cfg.initial_chunk_size_bytes;
    int max_dead_bytes_per_chunk = info.arena_cfg.max_dead_bytes_per_chunk == -1
                                       ? BFCArena::DEFAULT_MAX_DEAD_BYTES_PER_CHUNK
                                       : info.arena_cfg.max_dead_bytes_per_chunk;
    int initial_growth_chunk_size_bytes = info.arena_cfg.initial_growth_chunk_size_bytes == -1
                                              ? BFCArena::DEFAULT_INITIAL_GROWTH_CHUNK_SIZE_BYTES
                                              : info.arena_cfg.initial_growth_chunk_size_bytes;
    ArenaExtendStrategy arena_extend_str;
    switch (info.arena_cfg.arena_extend_strategy) {
      case static_cast<int>(ArenaExtendStrategy::kSameAsRequested):
        arena_extend_str = ArenaExtendStrategy::kSameAsRequested;
        break;
      case -1:  // default value supplied by user
      case static_cast<int>(ArenaExtendStrategy::kNextPowerOfTwo):
        arena_extend_str = ArenaExtendStrategy::kNextPowerOfTwo;
        break;
      default:
        LOGS_DEFAULT(ERROR) << "Received invalid value of arena_extend_strategy " << info.arena_cfg.arena_extend_strategy;
        return nullptr;
    }

#ifdef USE_MIMALLOC_ARENA_ALLOCATOR
    return std::shared_ptr<IAllocator>(
        std::make_unique<MiMallocAllocator>(max_mem));
#else
    return std::shared_ptr<IAllocator>(
        std::unique_ptr<BFCArena>(new BFCArena(std::move(device_allocator),
                                   max_mem,
                                   arena_extend_str,
                                   initial_chunk_size_bytes,
                                   max_dead_bytes_per_chunk,
                                   initial_growth_chunk_size_bytes)));
#endif
  }

  return AllocatorPtr(std::move(device_allocator));
}

#ifdef OPENGPU
AllocatorPtr &GPUAllocatorMgr::GetInstance()
{
    static GPUAllocatorMgr allocator;
    return allocator.allocator_ptr_;
}

GPUAllocatorMgr::GPUAllocatorMgr(){
    size_t avail;
    size_t total;
    cudaMemGetInfo( &avail, &total );
    bool gpu_create_arena = true;
    AllocatorCreationInfo device_info{[](int) { return std::unique_ptr<CUDAAllocator>(new CUDAAllocator(0, "CUDA")); },
                                        0, gpu_create_arena, {avail, 1, -1, -1, -1}};
    allocator_ptr_ = CreateAllocator(device_info);
};

GPUAllocatorMgr::~GPUAllocatorMgr(){};
#endif

AllocatorPtr &CPUAllocatorMgr::GetInstance()
{
    static CPUAllocatorMgr allocator;
    return allocator.allocator_ptr_;
}

CPUAllocatorMgr::CPUAllocatorMgr(){
    bool create_arena = true;
    AllocatorCreationInfo device_info{[](int) { return std::unique_ptr<TAllocator>(new TAllocator()); },
                                        0, create_arena, {0, 1, -1, -1, -1}};
    allocator_ptr_ = CreateAllocator(device_info);
};

CPUAllocatorMgr::~CPUAllocatorMgr(){};


}  // namespace bfcallocator
